// Tencent is pleased to support the open source community by making
// 蓝鲸智云 - 监控平台/日志平台 (BlueKing - Monitor/Log) available.
// Copyright (C) 2017-2022 THL A29 Limited, a Tencent company. All rights reserved.
// Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
// You may obtain a copy of the License at http://opensource.org/licenses/MIT
// Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
//

package oneapm

import (
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/TencentBlueKing/bkmonitor-kits/logger"
	oneapmproto "go.opentelemetry.io/collector/pkg/translator/oneapm/proto"
	"google.golang.org/protobuf/proto"

	"github.com/TencentBlueKing/bk-collector/define"
	"github.com/TencentBlueKing/bk-collector/receiver"
)

const (
	Source = "oneapm"

	routeTransactions   = "/Transactions"
	routeTransaction    = "/Transaction"
	routeAgentConnect   = "/Connect"
	routeAgentReconnect = "/Reconnect"
	routeAgentCommand   = "/AgentCommand"
	routeAgentData      = "/AgentData"
)

func init() {
	receiver.RegisterHttpRoute(Source, Encoders, []receiver.RouteWithFunc{
		{
			Method:       http.MethodPost,
			RelativePath: routeTransactions,
			HandlerFunc:  httpsrv.ExportTransactions,
		},
		{
			Method:       http.MethodPost,
			RelativePath: routeTransaction,
			HandlerFunc:  httpsrv.ExportTransaction,
		},
		{
			Method:       http.MethodPost,
			RelativePath: routeAgentData,
			HandlerFunc:  httpsrv.ExportAgentData,
		},
		{
			Method:       http.MethodPost,
			RelativePath: routeAgentConnect,
			HandlerFunc:  httpsrv.AgentConnect,
		},
		{
			Method:       http.MethodPost,
			RelativePath: routeAgentReconnect,
			HandlerFunc:  httpsrv.AgentConnect,
		},
		{
			Method:       http.MethodPost,
			RelativePath: routeAgentCommand,
			HandlerFunc:  httpsrv.AgentCommand,
		},
	})
}

type httpSrv struct{}

var httpsrv httpSrv

func (s httpSrv) export(w http.ResponseWriter, rh receiver.ResponseHandler, req *http.Request, rtype define.RecordType) {
	receiver.HttpExport(Source, w, rh, req, rtype)
}

func (s httpSrv) ExportTransactions(w http.ResponseWriter, req *http.Request) {
	responseHandler := s.getResponseHandler(req.Header.Get(receiver.ContentType), true)
	s.export(w, responseHandler, req, define.RecordTraces)
}

func (s httpSrv) ExportTransaction(w http.ResponseWriter, req *http.Request) {
	responseHandler := s.getResponseHandler(req.Header.Get(receiver.ContentType), false)
	s.export(w, responseHandler, req, define.RecordTraces)
}

func (s httpSrv) ExportAgentData(w http.ResponseWriter, req *http.Request) {
	responseHandler := s.getResponseHandler(req.Header.Get(receiver.ContentType), false)
	s.export(w, responseHandler, req, define.RecordMetrics)
}

func (s httpSrv) AgentCommand(w http.ResponseWriter, r *http.Request) {
	request := &oneapmproto.AgentCommandRequest{}
	if err := s.unmarshalRequest(r, request); err != nil {
		logger.Errorf("failed to unmarshal request: %v", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	response, err := s.handleAgentCommand(request)
	if err != nil {
		logger.Errorf("failed to handle request body: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	bs, err := s.marshalResponse(r, response)
	if err != nil {
		logger.Errorf("failed to marshal response: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	_, _ = w.Write(bs)
}

func (s httpSrv) AgentConnect(w http.ResponseWriter, r *http.Request) {
	request := &oneapmproto.AgentConnectRequest{}
	if err := s.unmarshalRequest(r, request); err != nil {
		logger.Errorf("failed to unmarshal request: %v", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	response := s.handleAgentConnect(request)
	bs, err := s.marshalResponse(r, response)
	if err != nil {
		logger.Errorf("failed to marshal response: %v", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	_, _ = w.Write(bs)
}

func (s httpSrv) unmarshalRequest(req *http.Request, msg proto.Message) error {
	contentType := req.Header.Get(receiver.ContentType)
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return err
	}
	defer req.Body.Close()

	switch contentType {
	case receiver.ContentTypeJson:
		err = json.Unmarshal(body, msg)
	default:
		err = proto.Unmarshal(body, msg)
	}
	return err
}

func (s httpSrv) marshalResponse(req *http.Request, msg proto.Message) ([]byte, error) {
	contentType := req.Header.Get(receiver.ContentType)
	var err error
	var bs []byte
	switch contentType {
	case receiver.ContentTypeJson:
		bs, err = json.Marshal(msg)
	default:
		bs, err = proto.Marshal(msg)
	}

	return bs, err
}

func (s httpSrv) getResponseHandler(contentType string, multi bool) receiver.ResponseHandler {
	switch contentType {
	case receiver.ContentTypeProtobuf:
		return HttpPbResponseHandler(multi)
	}
	// 缺省解析器为 contentTypeJson
	return HttpJsonResponseHandler(multi)
}

// HttpPb Response Handler

func HttpPbResponseHandler(multi bool) receiver.ResponseHandler {
	return httpPbResponseHandler{
		encoder: PbEncoder(multi),
	}
}

type httpPbResponseHandler struct {
	encoder receiver.Encoder
}

func (h httpPbResponseHandler) ContentType() string                              { return receiver.ContentTypeProtobuf }
func (h httpPbResponseHandler) Response(_ define.RecordType) ([]byte, error)     { return nil, nil }
func (h httpPbResponseHandler) MarshalErrorStatus(_ interface{}) ([]byte, error) { return nil, nil }

func (h httpPbResponseHandler) Unmarshal(rtype define.RecordType, b []byte) (define.RecordData, error) {
	return receiver.GetUnmarshal(h.encoder, rtype, b)
}

func (h httpPbResponseHandler) WriteResponse(w http.ResponseWriter, msg []byte) {
	receiver.WriteResponse(w, h.ContentType(), http.StatusOK, msg)
}

func (h httpPbResponseHandler) WriteError(w http.ResponseWriter, err error, statusCode int) {
	receiver.WriteError(w, h, err, statusCode)
}

// HttpJson Response Handler

func HttpJsonResponseHandler(multi bool) receiver.ResponseHandler {
	return httpJsonResponseHandler{
		encoder: JsonEncoder(multi),
	}
}

type httpJsonResponseHandler struct {
	encoder receiver.Encoder
}

func (h httpJsonResponseHandler) ContentType() string                              { return receiver.ContentTypeJson }
func (h httpJsonResponseHandler) Response(_ define.RecordType) ([]byte, error)     { return nil, nil }
func (h httpJsonResponseHandler) MarshalErrorStatus(_ interface{}) ([]byte, error) { return nil, nil }

func (h httpJsonResponseHandler) Unmarshal(rtype define.RecordType, b []byte) (define.RecordData, error) {
	return receiver.GetUnmarshal(h.encoder, rtype, b)
}

func (h httpJsonResponseHandler) WriteResponse(w http.ResponseWriter, msg []byte) {
	receiver.WriteResponse(w, h.ContentType(), http.StatusOK, msg)
}

func (h httpJsonResponseHandler) WriteError(w http.ResponseWriter, err error, statusCode int) {
	receiver.WriteError(w, h, err, statusCode)
}
